import random

import pygame
import numpy as np
from enum import Enum
from direct import *
import math
import os


class Maze:

    def __init__(self, name, rows, cols, start, end, symple_rate=0.1, reward_block_num=5):
        size = tuple((500, 500))
        self.screen = pygame.display.set_mode(size)
        self.screen.fill((255, 255, 255))
        self.rows = rows
        self.cols = cols
        self.blockW = size[0] // cols
        self.blockH = size[0] // rows
        self.start = start
        self.end = end
        self.sympleRate = symple_rate  # 简化率
        self.mazeFilePath = {'arr': f'mazeData/{name}_maze.npy', 'image': f'mazeData/{name}_maze.png',
                             'rwb_arr': f'mazeData/{name}_rwb.npy'}
        self.mazeImage = None
        # 用于记录一个block四边是否可走,0表示这个方向有墙
        self.maze = np.array([[[0 for i in range(4)] for j in range(cols)] for k in range(rows)])
        # 记录奖励点
        self.reward_block_num = reward_block_num
        self.reward_blocks = set()

    #  获得方格能采用的行动集
    def get_block_actions(self, state):
        return self.maze[state[0]][state[1]]

    def draw_maze(self):
        self.screen.fill((255, 255, 255))
        for i in range(self.rows):
            for j in range(self.cols):
                walls = np.where(self.get_block_actions((i, j)) == 0)[0]
                if walls.any():
                    # print(f'({i},{j}):{walls}')
                    for k in walls:
                        line = self.get_rect_border_position((i, j), DIRECTS[k], 0)
                        pygame.draw.line(self.screen, (0, 0, 0), line[0], line[1], 2)

    def generate_maze(self):
        map = np.zeros((self.rows, self.cols), dtype=int)
        rows, cols = self.rows, self.cols

        stack = [(0, 0)]  # 栈存储待访问的方块
        directions = DIRECTS  # 定义四个方向
        self.maze = np.array([[[0 for i in range(4)] for j in range(cols)] for k in range(rows)])
        while stack:
            current = stack.pop()
            x, y = current
            map[x][y] = 1  # 标记为访问过的方块
            neighbors = []
            for dx, dy in directions:
                new_x, new_y = x + dx, y + dy
                if 0 <= new_x < rows and 0 <= new_y < cols and map[new_x][new_y] == 0:
                    neighbors.append((new_x, new_y))
            if neighbors:
                stack.append(current)  # 当前方块重新压入栈
                next_block = neighbors[np.random.randint(len(neighbors))]  # 随机选择下一个方块
                dr = (next_block[0] - x, next_block[1] - y)  # 获取相对方向
                # line = self.get_rect_border_position((x, y), dr)  # 获取当前block与下一个block的边
                # pygame.draw.line(self.screen, (255, 255, 255), line[0], line[1], 3)  # 打破这条边
                # 更新迷宫
                dr_index = get_direct_index(dr)
                self.maze[x][y][dr_index] = 1
                self.maze[next_block[0]][next_block[1]][get_contrary_direct_index(dr)] = 1
                stack.append(next_block)
                # 演示过程
                # pygame.display.flip()
                # pygame.time.wait(100)

        # 简化迷宫，使得能有更多路径到达终点
        symple_num = int(rows * cols * self.sympleRate)
        for i in range(symple_num):
            # 随机在非外层中选择一个方格
            rand_x = random.randint(1, rows - 2)
            rand_y = random.randint(1, cols - 2)
            # 挑选一条未打破的边
            walls = np.where(self.maze[rand_x][rand_y] == 0)[0]
            if walls.any():
                dr_index = random.choice(walls)
                self.maze[rand_x][rand_y][dr_index] = 1
                dr = DIRECTS[dr_index]
                self.maze[rand_x+dr[0]][rand_y+dr[1]][get_contrary_direct_index(dr)] = 1

        # 查看迷宫是否正确
        # for i in range(len(self.mazeData)):
        #     for j in range(len(self.mazeData[i])):
        #         print(self.mazeData[i][j],end=' ')
        #     print()

        # 画出迷宫
        self.draw_maze()

        # 画出起点和终点位置
        self.draw_circle(self.start, (0, 255, 0))
        self.draw_circle(self.end, (0, 255, 0))

        # 保存迷宫
        np.save(self.mazeFilePath['arr'], arr=self.maze)
        pygame.image.save(self.screen, self.mazeFilePath['image'])
        self.mazeImage = pygame.image.load(self.mazeFilePath['image'])

        # 生成奖励点 会被更新，所以不能和迷宫一起保存
        self.generate_reward_blocks()
        self.draw_reward_block()
        np.save(self.mazeFilePath['rwb_arr'], arr=np.array(self.reward_blocks))


    def generate_reward_blocks(self):
        t = self.reward_block_num
        while t > 0:
            rand_x = random.randint(0, self.rows - 1)
            rand_y = random.randint(0, self.cols - 1)
            block = tuple((rand_x, rand_y))
            if block != self.start and block != self.end and block not in self.reward_blocks:
                t -= 1
                self.reward_blocks.add(block)

    def draw_reward_block(self):
        for block in self.reward_blocks:
            self.draw_circle(block, (255, 255, 0))

    # 从文件中加载地图
    def load_maze(self):
        if os.path.exists(self.mazeFilePath['arr']) and os.path.exists(self.mazeFilePath['image'] and os.path.exists(self.mazeFilePath['rwb_arr'])):
            self.maze = np.load(self.mazeFilePath['arr'])
            self.mazeImage = pygame.image.load(self.mazeFilePath['image'])
            self.reward_blocks = set(np.load(self.mazeFilePath['rwb_arr'],allow_pickle=True).tolist())
            self.screen.blit(self.mazeImage, (0, 0))
            self.draw_bot(self.start)
            self.draw_reward_block()
        else:
            self.generate_maze()

    # 重新绘制地图
    def reload_maze(self):
        self.screen.blit(self.mazeImage, (0, 0))
        self.draw_reward_block()

    # 绘制圆
    def draw_circle(self, position, color):
        w, h = self.blockW, self.blockH
        center = tuple((position[1] * w + w / 2, position[0] * h + h / 2))
        pygame.draw.circle(self.screen, color, center, 3 * (w + h) / 16)

    # 绘制机器人
    def draw_bot(self, position):
        self.draw_circle(position, (255, 0, 0))

    # 更新机器人位置
    def update_bot(self, position):
        # 移除奖励点
        if position in self.reward_blocks:
            self.reward_blocks.remove(position)
        self.reload_maze()
        self.draw_bot(position)
        pygame.display.flip()

    # 在两个相邻的各自间绘制线
    def draw_line_between_block(self, position1, position2, color):
        x1, y1 = position1[1], position1[0]
        x2, y2 = position2[1], position2[0]
        w, h = self.blockW, self.blockH
        start = (x1 * w + w / 2, y1 * h + h / 2)
        end = (x2 * w + w / 2, y2 * h + h / 2)
        pygame.draw.line(self.screen, color, start, end, 2)

    # 获得一个方格的某一条边
    def get_rect_border_position(self, position, direct, offset=1.5):
        x, y = position[1], position[0]
        w, h = self.blockW, self.blockH
        if direct == UP:
            return [(x * w + offset, y * h), ((x + 1) * w - offset, y * h)]
        elif direct == DOWN:
            return [(x * w + offset, (y + 1) * h), ((x + 1) * w - offset, (y + 1) * h)]
        elif direct == LEFT:
            return [(x * w, y * h + offset), (x * w, (y + 1) * h - offset)]
        elif direct == RIGHT:
            return [((x + 1) * w, y * h + offset), (x * w + w, (y + 1) * h - offset)]

    # 解迷宫
    def astart(self):
        print(1)
        rows, cols = self.rows, self.cols
        start = self.start
        end = self.end
        maze = self.maze
        open_set = {start}
        close_set = set()
        path_length = [[1000 for i in range(cols)] for j in range(rows)]  # 记录路径长度
        path_length[start[0]][start[1]] = 0
        path = [[(-1, -1) for i in range(cols)] for j in range(rows)]  # 记录路径
        f = lambda x: path_length[x[0]][x[1]] + math.fabs(end[0] - x[0]) + math.fabs(end[1] - x[1])  # 评价函数
        while open_set:
            # 找出优先级最高
            current = min(open_set, key=f)
            x, y = current
            if current == end:
                # 绘制解答路径
                # while path[x][y] != (-1, -1):
                #     self.draw_line_between_block((x, y), path[x][y], (0, 0, 255))
                #     pygame.time.wait(50)
                #     pygame.display.flip()
                #     x, y = path[x][y]
                break

            open_set.remove(current)
            close_set.add(current)
            for dr in DIRECTS:
                next = (x + dr[0], y + dr[1])
                dr_index = get_direct_index(dr)
                if 0 <= next[0] < rows and 0 <= next[1] < cols and maze[x][y][dr_index] == 1 and next not in close_set:
                    if next not in open_set:
                        path[next[0]][next[1]] = current
                        path_length[next[0]][next[1]] = path_length[current[0]][current[1]] + 1
                        self.draw_line_between_block(current, next, (255, 0, 0))
                        # self.update_bot(next)
                        open_set.add(next)
                    pygame.display.flip()
                    pygame.time.wait(50)


if __name__ == '__main__':
    # 初始化Pygame库
    pygame.init()
    # 设置标题和窗口大小
    pygame.display.set_caption("Maze Game")

    # 生成迷宫的二维数组
    maze = Maze('test', 10, 10, (0, 0), (9, 9), symple_rate=0.5)

    # 绘制迷宫
    # maze.generate_maze()
    maze.generate_maze()
    maze.astart()

    # 事件循环
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:  # 点击关闭按钮
                running = False
    # 退出Pygame库
    pygame.quit()
